#!/bin/bash
# 版本: 1.0
# 作者: huangsheng
# 兼容系统版本:ubuntu
# 验证:验证

# 检查操作系统是否为Ubuntu或Debian
OS_ID=$(grep '^ID=' /etc/os-release | cut -d= -f2 | tr -d '\"')
if [[ ! "$OS_ID" =~ ^(ubuntu|debian)$ ]]; then
    echo "此脚本仅支持Ubuntu或Debian系统。"
    exit 1
fi

# 用户选择要安装的ROCm版本
echo "请选择要安装的ROCm版本:"
echo "0) ROCm 6.1.3"
echo "1) ROCm 6.2"
echo "2) ROCm 6.3"
echo "3) ROCm 6.4"
read -p "输入选择 (0，1, 2 或 3): " version_choice

if [ "$version_choice" = "0" ]; then
    ROCM_VER="6.1.3"
    ROCM_PACK_URL="https://repo.radeon.com/amdgpu-install/6.1.3/ubuntu/jammy/amdgpu-install_6.1.60103-1_all.deb"
elif [ "$version_choice" = "1" ]; then
    ROCM_VER="6.2"
    ROCM_PACK_URL="https://repo.radeon.com/amdgpu-install/6.2.4/ubuntu/jammy/amdgpu-install_6.2.60204-1_all.deb"
elif [ "$version_choice" = "2" ]; then
    ROCM_VER="6.3"
    ROCM_PACK_URL="https://repo.radeon.com/amdgpu-install/6.3.4/ubuntu/jammy/amdgpu-install_6.3.60304-1_all.deb"
elif [ "$version_choice" = "3" ]; then
    ROCM_VER="6.4"
    ROCM_PACK_URL="https://repo.radeon.com/amdgpu-install/6.4/ubuntu/jammy/amdgpu-install_6.4.60400-1_all.deb"
else
    echo "无效选择，脚本退出。"
    exit 1
fi

# 更新系统并安装必要的软件包
sudo apt update && sudo apt upgrade -y
sudo apt install -y git python3-pip python3-venv python3-dev libstdc++-12-dev wget curl cmake

echo "开始安装ROCm ${ROCM_VER} ..."

# 下载并安装 ROCm
if [ ! -f "$(basename ${ROCM_PACK_URL})" ]; then
    wget $ROCM_PACK_URL
fi
sudo dpkg -i $(basename ${ROCM_PACK_URL})

sudo amdgpu-install -y --usecase=graphics,rocm

# 安装 ROCm 开发包
sudo apt install -y rocm-hip-runtime rocm-hip-sdk rocm-ml-libraries
sudo apt install -y rocm-libs rocm-dev rocm-device-libs rocm-utils

# 设置环境变量
sudo ln -s /opt/rocm/bin/rocm-smi /usr/local/bin/rocm-smi
sudo usermod -a -G render,video $LOGNAME

#支持多显卡
#grep -q 'GRUB_CMDLINE_LINUX="iommu=pt"' /etc/default/grub || sudo sed -i '/^GRUB_CMDLINE_LINUX=/c\GRUB_CMDLINE_LINUX="iommu=pt"' /etc/default/grub
#sudo update-grub

# 安装完成后的提示信息  
echo "ROCm 安装完成。请重启系统或重新加载内核模块以确保一切正常。"
